from keras.datasets import cifar10
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential, Graph
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import Convolution2D, MaxPooling2D,UpSampling2D
from keras.optimizers import SGD
from keras.utils import np_utils
import keras.backend as K
from keras.regularizers import *

def custom_loss(y_true,y_pred):
    MSE = K.mean(K.square(y_pred - y_true))
    return MSE

def baseline_conv(img_channels,img_rows,img_cols):
    model = Sequential()
    model.add(Convolution2D(32, 3, 3, b_constraint=0,activation='tanh',
                            input_shape=(img_channels, img_rows, img_cols)))
    model.add(Convolution2D(32, 3, 3, b_constraint=0,activation='tanh',))
    model.add(Convolution2D(1, 3, 3, b_constraint=0))
    model.summary()
    return model

def baseline_conv_test(img_channels,img_rows,img_cols):
    model = Sequential()
    model.add(Convolution2D(32, 3, 3, border_mode='same',  b_constraint=0,activation='tanh',
                            input_shape=(img_channels, img_rows, img_cols)))
    model.add(Convolution2D(32, 3, 3, border_mode='same', b_constraint=0,activation='tanh'))
    model.add(Convolution2D(1, 3, 3, border_mode='same', b_constraint=0))
    model.summary()
    return model

def up_pooling_model(img_channels,img_rows,img_cols):
    model = Sequential()
    model.add(Convolution2D(32, 3, 3, border_mode='same', b_constraint=0,
                            input_shape=(img_channels, img_rows, img_cols)))
    model.add(Convolution2D(64, 3, 3, border_mode='same', b_constraint=0))
    model.add(MaxPooling2D(pool_size=(2,2)))
    model.add(Dropout(0.25))

    model.add(UpSampling2D(size=(2,2)))
    model.add(Convolution2D(64,1, 1, b_constraint=0))
    model.add(Convolution2D(32,3, 3, border_mode='same', b_constraint=0))
    model.add(Dropout(0.25))

    model.add(Convolution2D(1,3,3,border_mode='same', b_constraint=0))
    model.summary()
    return model

def linear_conv(img_channels,img_rows,img_cols):
    model = Sequential()
    model.add(Convolution2D(1, 5, 5, b_constraint=0,
                            input_shape=(img_channels, img_rows, img_cols)))
    model.add(Convolution2D(1, 5, 5, b_constraint=0))
    model.summary()
    return model

def linear_conv_test(img_channels,img_rows,img_cols):
    model = Sequential()
    model.add(Convolution2D(1, 5, 5, b_constraint=0,border_mode='same',
                            input_shape=(img_channels, img_rows, img_cols)))
    model.add(Convolution2D(1, 5, 5, b_constraint=0,border_mode='same'))
    model.summary()
    return model

def SRCNN(img_channels,img_rows,img_cols):
    model = Sequential()
    model.add(Convolution2D(64, 9, 9, b_constraint=0,activation='relu',border_mode='same',
                            input_shape=(img_channels, img_rows, img_cols)))
    model.add(Convolution2D(32, 1, 1, activation='relu',b_constraint=0))
    model.add(Convolution2D(img_channels, 5, 5, b_constraint=0,border_mode='same'))
    model.summary()
    return model

def SRCNN_test(img_channels,img_rows,img_cols):
    model = Sequential()
    model.add(Convolution2D(64, 9, 9, border_mode='same',b_constraint=0,activation='relu',
                            input_shape=(img_channels, img_rows, img_cols)))
    model.add(Convolution2D(32, 1, 1, b_constraint=0,activation='relu'))
    model.add(Convolution2D(img_channels, 5, 5, border_mode='same',b_constraint=0))
    model.summary()
    return model

def residual_conv(img_channels,img_rows,img_cols):
    conv = Sequential()
    conv.add(Convolution2D(32, 3, 3, b_constraint=0,activation='tanh',border_mode='same',
                            input_shape=(img_channels, img_rows, img_cols)))
    conv.add(Convolution2D(32, 3, 3, b_constraint=0,activation='tanh',border_mode='same'))
    conv.add(Convolution2D(1, 3, 3, b_constraint=0,border_mode='same'))

    skip = Sequential()
    skip.add(Activation(activation='linear',input_shape=(img_channels, img_rows, img_cols)))

    model = Graph()
    model.add_input(name='conv_input',input_shape=(img_channels, img_rows, img_cols))
    model.add_input(name='skip_input',input_shape=(img_channels, img_rows, img_cols))
    model.add_node(conv,name='conv',input='conv_input')
    model.add_node(skip,name='skip',input='skip_input')
    model.add_output(name='output',inputs=['conv','skip'],merge_mode='sum')
    return model

def residual_conv_test(img_channels,img_rows,img_cols):
    conv = Sequential()
    conv.add(Convolution2D(32, 3, 3, b_constraint=0,activation='tanh',border_mode='same',
                            input_shape=(img_channels, img_rows, img_cols)))
    conv.add(Convolution2D(32, 3, 3, b_constraint=0,activation='tanh',border_mode='same'))
    conv.add(Convolution2D(1, 3, 3, b_constraint=0,border_mode='same'))

    skip = Sequential()
    skip.add(Activation(activation='linear',input_shape=(img_channels, img_rows, img_cols)))

    model = Graph()
    model.add_input(name='conv_input',input_shape=(img_channels, img_rows, img_cols))
    model.add_input(name='skip_input',input_shape=(img_channels, img_rows, img_cols))
    model.add_node(conv,name='conv',input='conv_input')
    model.add_node(skip,name='skip',input='skip_input')
    model.add_output(name='output',inputs=['conv','skip'],merge_mode='sum')
    return model